# -*- coding: utf-8 -*-
import torch
from cubetools import image_tools
from cubetools.video_predict import VideoPredict
from cubetools.huggingface import download_model
from ultralyticsplus import YOLO, render_result


class ModelPredict(object):
    def __init__(self):
        self.video_predict = VideoPredict(callback_predict=self.predict)

        model_path = download_model('keremberke/yolov8s-pothole-segmentation', exclude=[])
        self.model = YOLO(model_path + '/best.pt')
        # set model parameters
        self.model.overrides['conf'] = 0.25  # NMS confidence threshold
        self.model.overrides['iou'] = 0.45  # NMS IoU threshold
        self.model.overrides['agnostic_nms'] = False  # NMS class-agnostic
        self.model.overrides['max_det'] = 1000  # maximum number of detections per image

    def predict(self, img):
        img_pil, img = image_tools.read_img(img)

        with torch.no_grad():
            results = self.model.predict(img_pil)

        img_pil = render_result(model=self.model, image=img_pil, result=results[0])
        img_url = image_tools.pil2url(img_pil)
        return [], img_url

    def predict_video(self, url):
        return self.video_predict.read_predict(url)[1]
